#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
Coordinate calculator for target selection and prediction.
Ported from legacy src/utils/coordinate.py with configuration injection.
"""

import logging
import time
import uuid
from typing import Dict, List, Optional, Tuple

import numpy as np

from ...config.settings import TrackingConfig
from .kalman_filter import KalmanFilter

logger = logging.getLogger("RDK_YOLO")


class CoordinateCalculator:
    """
    负责多目标评分、卡尔曼预测和坐标换算的辅助类。
    """

    def __init__(self, frame_width: int, frame_height: int, tracking_config: TrackingConfig):
        self.frame_width = frame_width
        self.frame_height = frame_height
        self.center_x = frame_width // 2
        self.center_y = frame_height // 2
        self.latest_bbox: Optional[Tuple[int, int, int, int]] = None
        self.latest_update_time: float = 0.0

        self.tracking_config = tracking_config
        target_cfg = tracking_config.target_tracking

        # 评分权重
        self.confidence_weight = target_cfg.confidence_weight
        self.size_weight = target_cfg.size_weight
        self.persistence_weight = target_cfg.persistence_weight
        self.position_weight = target_cfg.position_weight

        # 目标跟踪参数
        self.target_switch_threshold = target_cfg.target_switch_threshold
        self.target_timeout = target_cfg.target_timeout
        self.max_history_count = target_cfg.max_history_count
        self.max_position_distance = target_cfg.max_position_distance

        # 方向判断
        self.direction_threshold = target_cfg.direction_threshold

        # 预测开关
        self.prediction_enabled = tracking_config.prediction_enabled
        self.last_predicted_position: Optional[Tuple[float, float]] = None
        self.prediction_confidence: float = 0.0

        # 历史状态
        self.current_target_id: Optional[str] = None
        self.target_history: Dict[str, Dict[str, any]] = {}

        # 卡尔曼滤波配置
        kalman_cfg = tracking_config.kalman
        self.kalman_dt = kalman_cfg.dt
        self.process_noise_scale = kalman_cfg.process_noise_scale
        self.measurement_noise_scale = kalman_cfg.measurement_noise_scale
        self.kalman_filters: Dict[str, KalmanFilter] = {}

    # 核心 API ---------------------------------------------------------------

    def update_targets(self, bboxes: List[Tuple[int, int, int, int]], scores: List[float]):
        """
        结合检测框和置信度选择当前追踪目标。
        返回 (target_id, bbox)。
        """
        current_time = time.time()

        if len(bboxes) == 0:
            return self._handle_no_detections(current_time)

        self._cleanup_old_targets(current_time)

        centers = []
        areas = []
        for bbox in bboxes:
            x1, y1, x2, y2 = bbox
            centers.append(((x1 + x2) // 2, (y1 + y2) // 2))
            areas.append((x2 - x1) * (y2 - y1))

        target_ids = self._associate_targets(centers, areas, bboxes, scores, current_time)

        normalized_areas = [area / max(areas) if areas else 0 for area in areas]

        combined_scores = []
        persistence_scores = []
        position_scores = []

        for i, target_id in enumerate(target_ids):
            history = self.target_history[target_id]

            persistence_score = history["count"] / self.max_history_count
            persistence_scores.append(persistence_score)

            position_score = history.get("position_score", 0.0)
            position_scores.append(position_score)

            base_score = (self.confidence_weight * scores[i] +
                          self.size_weight * normalized_areas[i])

            current_target_bonus = self.target_switch_threshold if self.current_target_id == target_id else 0.0

            final_score = (base_score +
                           persistence_score * self.persistence_weight +
                           position_score * self.position_weight +
                           current_target_bonus)

            combined_scores.append(final_score)

        if combined_scores:
            best_idx = self._select_best_target(combined_scores, target_ids)
            best_id = target_ids[best_idx]
            self.current_target_id = best_id
            self.latest_bbox = bboxes[best_idx]
            self.latest_update_time = current_time

            if self.prediction_enabled:
                self._update_prediction(best_id, centers[best_idx], current_time)
                self.prediction_confidence = 1.0

            return best_id, self.latest_bbox

        # 后备策略
        best_idx = int(np.argmax(scores)) if scores else 0
        best_id = target_ids[best_idx]
        self.current_target_id = best_id
        self.latest_bbox = bboxes[best_idx]
        self.latest_update_time = current_time

        if self.prediction_enabled:
            self._update_prediction(best_id, centers[best_idx], current_time)
            self.prediction_confidence = 1.0

        return best_id, self.latest_bbox

    # 预测与状态 ------------------------------------------------------------

    def _handle_no_detections(self, current_time):
        if self.current_target_id is not None and self.prediction_enabled:
            predicted = self._predict_target_position(self.current_target_id, current_time)
            if predicted is not None:
                logger.debug("目标丢失但使用预测位置: ID=%s, 预测边界框=%s",
                             self.current_target_id, predicted)
                self.latest_bbox = predicted
                self.latest_update_time = current_time
                self.prediction_confidence -= 0.2

                if self.prediction_confidence <= 0:
                    self._reset_prediction_state()
                    return None, None

                return self.current_target_id, self.latest_bbox

        self._reset_prediction_state()
        return None, None

    def _reset_prediction_state(self):
        self.current_target_id = None
        self.latest_bbox = None
        self.latest_update_time = 0.0
        self.last_predicted_position = None
        self.prediction_confidence = 0.0

    def _update_prediction(self, target_id: str, center: Tuple[int, int], current_time: float) -> None:
        kf = self.kalman_filters.get(target_id)
        if kf is None:
            kf = KalmanFilter(
                dt=self.kalman_dt,
                process_noise_scale=self.process_noise_scale,
                measurement_noise_scale=self.measurement_noise_scale,
            )
            self.kalman_filters[target_id] = kf

        estimated_position, estimated_velocity = kf.update(center, current_time)

        if target_id in self.target_history:
            self.target_history[target_id]["velocity"] = estimated_velocity
            velocity_magnitude = np.sqrt(estimated_velocity[0] ** 2 + estimated_velocity[1] ** 2)
            self.target_history[target_id]["velocity_magnitude"] = velocity_magnitude
            logger.debug(
                "目标速度更新: ID=%s, 速度=(%.1f, %.1f), 大小=%.1f",
                target_id,
                estimated_velocity[0],
                estimated_velocity[1],
                velocity_magnitude,
            )

    def _predict_target_position(self, target_id: str, current_time: float):
        if (target_id not in self.target_history or
                target_id not in self.kalman_filters or
                not self.kalman_filters[target_id].initialized):
            return None

        history = self.target_history[target_id]
        last_seen = history["last_seen"]
        last_bbox = history["bbox"]

        if current_time - last_seen > self.target_timeout * 1.5:
            return None

        prediction_time = current_time - last_seen
        predicted_position = self.kalman_filters[target_id].predict(prediction_time)

        if predicted_position is None:
            return None

        self.last_predicted_position = predicted_position

        x1, y1, x2, y2 = last_bbox
        width = x2 - x1
        height = y2 - y1

        pred_x1 = int(predicted_position[0] - width / 2)
        pred_y1 = int(predicted_position[1] - height / 2)
        pred_x2 = int(predicted_position[0] + width / 2)
        pred_y2 = int(predicted_position[1] + height / 2)

        pred_x1 = max(0, pred_x1)
        pred_y1 = max(0, pred_y1)
        pred_x2 = min(self.frame_width, pred_x2)
        pred_y2 = min(self.frame_height, pred_y2)

        if pred_x2 <= pred_x1 or pred_y2 <= pred_y1:
            return None

        predicted_bbox = (pred_x1, pred_y1, pred_x2, pred_y2)
        logger.debug(
            "目标位置预测: ID=%s, 原始=%s, 预测=%s, 预测时间=%.3f秒",
            target_id,
            last_bbox,
            predicted_bbox,
            prediction_time,
        )
        return predicted_bbox

    # 关联与维护 ------------------------------------------------------------

    def _associate_targets(
        self,
        centers: List[Tuple[int, int]],
        areas: List[int],
        bboxes: List[Tuple[int, int, int, int]],
        scores: List[float],
        current_time: float,
    ) -> List[str]:
        target_ids: List[str] = []

        if not self.target_history:
            for center, area, bbox, score in zip(centers, areas, bboxes, scores):
                target_id = f"target_{uuid.uuid4().hex[:8]}"
                self.target_history[target_id] = {
                    "count": 1,
                    "last_seen": current_time,
                    "position": center,
                    "bbox": bbox,
                    "score": score,
                    "area": area,
                    "position_score": 0.0,
                    "velocity": (0.0, 0.0),
                    "velocity_magnitude": 0.0,
                }
                target_ids.append(target_id)
            return target_ids

        history_centers = [info["position"] for info in self.target_history.values()]
        history_ids = list(self.target_history.keys())

        distance_matrix = np.zeros((len(centers), len(history_centers)))
        for i, current_center in enumerate(centers):
            for j, history_center in enumerate(history_centers):
                distance = np.sqrt(
                    (current_center[0] - history_center[0]) ** 2 +
                    (current_center[1] - history_center[1]) ** 2
                )
                distance_matrix[i, j] = distance

        assigned_history_indices: List[int] = []

        for i, (current_center, current_area, current_bbox, current_score) in enumerate(
            zip(centers, areas, bboxes, scores)
        ):
            min_distance_idx = -1
            min_distance = float("inf")

            for j in range(len(history_centers)):
                if j in assigned_history_indices:
                    continue

                distance = distance_matrix[i, j]
                target_id = history_ids[j]

                if self.prediction_enabled and target_id in self.kalman_filters and self.kalman_filters[target_id].initialized:
                    last_position = self.target_history[target_id]["position"]
                    last_time = self.target_history[target_id]["last_seen"]
                    dt = current_time - last_time
                    predicted_position = self.kalman_filters[target_id].predict(dt)

                    if predicted_position is not None:
                        pred_distance = np.sqrt(
                            (current_center[0] - predicted_position[0]) ** 2 +
                            (current_center[1] - predicted_position[1]) ** 2
                        )
                        distance = distance * 0.3 + pred_distance * 0.7

                if distance > self.max_position_distance:
                    continue

                if distance < min_distance:
                    min_distance = distance
                    min_distance_idx = j

            if min_distance_idx == -1:
                target_id = f"target_{uuid.uuid4().hex[:8]}"
                self.target_history[target_id] = {
                    "count": 1,
                    "last_seen": current_time,
                    "position": current_center,
                    "bbox": current_bbox,
                    "score": current_score,
                    "area": current_area,
                    "position_score": 0.0,
                    "velocity": (0.0, 0.0),
                    "velocity_magnitude": 0.0,
                }
                target_ids.append(target_id)
            else:
                assigned_history_indices.append(min_distance_idx)
                target_id = history_ids[min_distance_idx]
                self._update_history_entry(target_id, current_center, current_bbox, current_score, current_area, current_time)
                target_ids.append(target_id)

        return target_ids

    def _update_history_entry(
        self,
        target_id: str,
        center: Tuple[int, int],
        bbox: Tuple[int, int, int, int],
        score: float,
        area: int,
        current_time: float,
    ) -> None:
        history = self.target_history[target_id]
        history["count"] = min(history["count"] + 1, self.max_history_count)
        history["last_seen"] = current_time
        history["position"] = center
        history["bbox"] = bbox
        history["score"] = score
        history["area"] = area

        if self.current_target_id == target_id and history.get("velocity") is not None:
            velocity = history["velocity"]
            velocity_magnitude = history.get("velocity_magnitude", 0.0)
            logger.debug(
                "目标运动: ID=%s, 速度=(%.1f, %.1f), 速度大小=%.1f",
                target_id,
                velocity[0],
                velocity[1],
                velocity_magnitude,
            )

    def _cleanup_old_targets(self, current_time: float) -> None:
        to_remove = []
        for target_id, info in self.target_history.items():
            if current_time - info["last_seen"] > self.target_timeout:
                to_remove.append(target_id)

        for target_id in to_remove:
            logger.debug("清理过期目标: %s", target_id)
            self.target_history.pop(target_id, None)
            self.kalman_filters.pop(target_id, None)

        if self.current_target_id in to_remove:
            self.current_target_id = None
            self.latest_bbox = None
            self.latest_update_time = 0.0

    def _select_best_target(self, combined_scores: List[float], target_ids: List[str]) -> int:
        best_idx = int(np.argmax(combined_scores))
        best_target_id = target_ids[best_idx]

        if (self.current_target_id is not None and
                self.current_target_id in target_ids and
                best_target_id != self.current_target_id):
            current_idx = target_ids.index(self.current_target_id)
            current_score = combined_scores[current_idx]
            best_score = combined_scores[best_idx]

            if best_score - current_score < self.target_switch_threshold:
                return current_idx

        return best_idx

    # ------------------------------------------------------------------ #
    # Utility helpers

    def calculate_offset(self, bbox: Tuple[int, int, int, int]) -> Tuple[float, float, float]:
        if bbox is None:
            return 0.0, 0.0, 0.0

        x1, y1, x2, y2 = bbox
        target_center_x = (x1 + x2) // 2
        target_center_y = (y1 + y2) // 2

        offset_x = target_center_x - self.center_x
        offset_y = target_center_y - self.center_y
        distance = np.hypot(offset_x, offset_y)

        return float(offset_x), float(offset_y), float(distance)

    def get_direction(self, offset_x: float, offset_y: float) -> str:
        threshold = self.direction_threshold
        if abs(offset_x) < threshold and abs(offset_y) < threshold:
            return "center"

        directions = []
        if offset_x > threshold:
            directions.append("right")
        elif offset_x < -threshold:
            directions.append("left")

        if offset_y > threshold:
            directions.append("down")
        elif offset_y < -threshold:
            directions.append("up")

        return "-".join(directions)
